/**
 * 
 */
package mdp.solve;

import java.util.ArrayList;
import java.util.Collections;
import java.util.Comparator;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ArrayBlockingQueue;

import RC.structs.DTree;
import RC.structs.DTreeNode;

import util.Pair;
import util.Timer;

import dd.discrete.DD;
import mdp.id.ADDHandler;
import mdp.id.RDDL2DD;
import mdp.id.RDDL2DD.DEBUG;
import mdp.id.RDDL2DD.ORDER;

/**
 * @author NR
 *
 * TODO
 */
public class MaxDecomposition<T extends RDDL2DD>  implements Runnable{
	
	private static double	EPSILON	= 0;
	private T _handler;
	private ArrayBlockingQueue<Pair<T, Integer>> _bq = null;
	
	private static final boolean	MAX_ON	= true;
	/**
	 * 
	 */
	private ArrayList<String>	_maxOrder;
	private ArrayList<String>	_sumOrder;
	private ArrayList<Integer>	_rewards = null;
	
	public MaxDecomposition(T handler, String domain, String instance, double epsilon,
			ArrayBlockingQueue<Pair<T, Integer>> bq) {
		_bq = bq; 
		EPSILON = epsilon;
		_handler = handler;
		_handler.buildFromFile(domain, instance);
		_sumOrder = handler.getSumOrderString();
		_maxOrder = handler.getMaxOrderString();
		
//		_dtree = new DTree<T>(handler, new HashSet<String>(_sumOrder), _maxOrder );
	}

	/* (non-Javadoc)
	 * @see java.lang.Runnable#run()
	 */
//	@Override
	public void run() {
		
//		if( _handler._bRewardActionDependent ) {
//			System.err.println("this procedure only applicable for action independent rewards");
//			System.exit(1);
//		}
		
		int _valueDD = _handler.DD_ZERO;
		int iter = 0;
		boolean done = false;
		Timer solutionTimer = new Timer();
		
		while( !done ) {
			solutionTimer.ResumeTimer();
			
			int Vprev = _valueDD;
			int Vprimed = _handler.swapForPrimes(_valueDD);
			_handler.saveThis(Vprev);
			_valueDD = maxRegress(Vprimed, _valueDD);
			
			int diff = _handler._context.applyInt(_valueDD, Vprev, DD.ARITH_MINUS);
			
			double error = Math.max(
					Math.abs(_handler.getMax(diff)), 
					Math.abs(_handler.getMin(diff) == _handler.getNegInf() ? 0 : _handler.getMin(diff)));
			++iter;
			
			if( error < EPSILON || _handler.isHorizon(iter+1) ) {
				done = true;
			}
			
			solutionTimer.PauseTimer();
			
			System.out.println("Iteration " + iter + " " + error + " " 
					+ solutionTimer.GetElapsedTimeInMinutes() + " " + 
					_handler.getSize(_valueDD));
			_handler.forgetThis(Vprev);
			_handler.flushCaches(true, _valueDD);
		}
		
		solutionTimer.ResumeTimer();
		
		int Vprimed = _handler.swapForPrimes(_valueDD);
		Pair<Integer, Integer> vq = _handler.regress(Vprimed, null, true);
		
		_valueDD = vq._o1;
//		int _policyDD = _handler.makePolicy(vq._o2);
		int _policy = _handler.makePolicy(vq._o1, vq._o2);
//		_valueDD = vq._o1;
//		_valueDD = _handler.fix(_valueDD);

		solutionTimer.PauseTimer();
		
		_handler.setValueDD(_valueDD);
		_handler.setPolicyDD(_policy);
		
		ArrayList stats = new ArrayList();
		
		try{
			stats = _handler.playDD(_policy);
		}catch(Exception e) {
			e.printStackTrace();
		}
		
		stats.add(0);
		stats.add(solutionTimer.GetElapsedTimeInMinutes());
		stats.add(_handler.getSize(_valueDD));
		stats.add(_handler.getSize(_policy));
		
		System.out.println(stats);
//		_handler.showGraph(_policyDD);
		
		if( _bq != null ) {
			try
			{
				_bq.put(new Pair<T, Integer>(_handler, _valueDD));
//				_bq.put(new Pair<T, Integer>(_handler, _policy));
			} catch (InterruptedException e)
			{
				e.printStackTrace();
			}
		
		}
	}
	
	/**
	 * @param pVprimed
	 * @param pObject
	 * @return
	 */
	private int maxRegress(int pVprimed, final int... save) {
		
		Pair<Integer, Integer> ret = _handler.regressExogenous(pVprimed, null, true, save);

		int vp = ret._o1;
		int sumInd = ret._o2;
		
		int vx = regressControl(vp, sumInd);
		
		return vx;
		
	}

	/**
	 * @param pVp
	 * @param pSumInd
	 * @param pMaxInd
	 * @param pHashMap
	 * @param pRoot
	 * @param pHorizon
	 * @return
	 */
	private int regressControl(int pVp, int pSumInd) {

		int Vcur = _handler.ConstrainThis(pVp, null, true, true);
		int maxInd = 0;
		
		ArrayList<Pair<Integer, String>> scores = _handler.getScores(
				new HashMap<Integer, Boolean>(), _sumOrder.size());
		
//		System.out.println(scores);
//		
		HashMap<Integer, Boolean> maxed = new HashMap<Integer, Boolean>();
		
		for( int i = pSumInd; i < _sumOrder.size(); ++i ) {

				String var = _sumOrder.get(i);
				int q = _handler.computeExpectation(Vcur, _handler.toID(var), null, true);

				updateScores(scores, i);
				
				if( _handler._debug.compareTo(DEBUG.SOLUTION_INFO) >= 0 ) {
					System.out.println("scores : " + scores);
				}
				
				
				Vcur = maxOutZeros(q, scores, maxed);
				
//				if( makePolicy ) {
//					int qfixed = _handler.fix(_handler._context.scalarAdd(q, 1.0d));
//					int vfixed = _handler.fix(Vcur);
//					
//					int diff = _handler._context.applyInt(vfixed, qfixed, DD.ARITH_MINUS);
//					int actzero = _handler.thresholdAt(diff, -1.0d);
//					int actone = _handler._context.applyInt(_handler.DD_ONE, actzero, DD.ARITH_MINUS);
//					
//					_handler.forgetThis(policy);
//					
//					policy = _handler._context.applyInt(policy,	
//						actone,
//						DD.ARITH_PROD);
					
//				}
				
		}
		
//		Vcur = _handler._context.scalarMultiply(Vcur, _handler.getDiscount());
		
		for( int r : _rewards ) {
			if( _handler._debug.compareTo(DEBUG.SOLUTION_INFO) >= 0 ) {
				System.out.println("adding reward " + r );
			}
			
			Vcur = _handler._context.applyInt(Vcur, r, DD.ARITH_SUM);
		}
		
		Vcur = _handler.fix(Vcur);
		
		_rewards = null;
		
		return Vcur;
	}

	
	/**
	 * @param pQ
	 * @param pScores
	 * @param pMaxed 
	 * @return
	 */
	private int maxOutZeros(int pQ, ArrayList<Pair<Integer, String>> pScores, HashMap<Integer, Boolean> pMaxed) {
		int ret = pQ;
		ArrayList<Pair<Integer, String>> dels = new ArrayList<Pair<Integer,String>>();
		
		for( Pair<Integer, String> p : pScores ) {
			if( p._o1 == 0 ) {
				String act = p._o2;

				int id = _handler.toContextID(act);
				
				ArrayList<Integer> rew = getActionReward(act);
				
				for( int r : rew ) {
					if( _handler._debug.compareTo(DEBUG.SOLUTION_INFO) >= 0 ) {
						System.out.println("Adding reward : " + r  );
					}
					
					ret = _handler._context.applyInt(ret, r, DD.ARITH_SUM);	
				}
				
				if( _handler._debug.compareTo(DEBUG.SOLUTION_INFO) >= 0 ) {
					System.out.println("Maxing : " + act  );
				}
				
				ret = _handler._context.opOut(ret, id, DD.ARITH_MAX);
				dels.add(p);
				pMaxed.put(_handler.toID(act), false);
			}
		}
		
		for( Pair<Integer, String> p : dels ) {
			pScores.remove(p);
		}
		
		return ret;
		
	}

	/**
	 * @param pId
	 * @return
	 */
	private ArrayList<Integer> getActionReward(String act) {
		if( _rewards == null ) {
			_rewards = new ArrayList<Integer>(_handler.getRewards());
		}
		
		ArrayList<Integer> ret = new ArrayList<Integer>();
		
		for( int r : _rewards ) {
			Set gids = _handler._context.getGIDs(r);
			if( gids.contains(_handler.toContextID(act)) ) {
				ret.add(r);
			}
		}
		
		for( Integer r : ret ) {
			_rewards.remove(r);
		}
		
		return ret;
		
	}

	/**
	 * @param pScores
	 * @param pI
	 * @return
	 */
	private void updateScores(
			ArrayList<Pair<Integer, String>> pScores, int pI) {
		String xp = _sumOrder.get(pI);
		ArrayList<String> acts = (ArrayList<String>) _handler._hmvars2act.get(xp);
		for( Pair<Integer, String> p : pScores ) {
			if( acts.contains(p._o2) ) {
				--p._o1;
			}
		}
	}

	/**
	 * @param pCutSet
	 * @return
	 */
	private ArrayList<String> retainActionsOnly(Set<String> pCutSet) {
		Set<String> copy = new HashSet<String>(pCutSet);
		copy.retainAll(_handler._alactionvars);
		return new ArrayList<String>(copy);
	}

	/**
	 * @return
	 */
	

	public static void main(String[] args) throws InterruptedException {
		ADDHandler handle = new ADDHandler(DEBUG.DIAGRAMS , ORDER.GUESS );
		Runnable vic = 
			new MaxDecomposition<ADDHandler>(handle, args[0], args[1], 
					Double.parseDouble(args[2]), null);
		Thread vicWorker = new Thread(vic, "VIC");
		vicWorker.start();
		vicWorker.join();
	}
	
}
